# 定义评估函数
# output=AI调用的结果，expected_output=langfuse中期望输出的记过
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langfuse import Langfuse
# 加载 .env 到环境变量
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())


def simple_evaluation(output, expected_output):
    return output == expected_output


langfuse = Langfuse()

need_answer = PromptTemplate.from_template("""
*********
你是AIGC课程的助教，你的工作是从学员的课堂交流中选择出需要老师回答的问题，加以整理以交给老师回答。

课程内容:
{outlines}
*********
学员输入:
{user_input}
*********
如果这是一个需要老师答疑的问题，回复Y，否则回复N。
只回复Y或N，不要回复其他内容。""")

model = ChatOpenAI(temperature=0, model_kwargs={"seed": 42})
parser = StrOutputParser()
chain_v1 = (
        need_answer
        | model
        | parser
)


# 创建执行评估的函数
def run_evaluation(chain, dataset_name, run_name):
    dataset = langfuse.get_dataset(dataset_name)

    def process_item(item):
        handler = item.get_langchain_handler(run_name=run_name)
        output = chain.invoke(item.input, config={"callbacks": [handler]})
        # Assuming handler.root_span.score is a synchronous function 评估数据集
        handler.trace.score(
            name="accuracy",
            value=simple_evaluation(output, item.expected_output)
        )
        print('.', end='', flush=True)

    # for item in dataset.items:
    #     process_item(item)
    # 建议并行处理
    with ThreadPoolExecutor(max_workers=8) as executor:
        executor.map(process_item, dataset.items)


run_evaluation(chain_v1, "my-dataset", "v1-" + datetime.now().strftime("%d/%m/%Y %H:%M:%S"))
'''
评估完成以后，在langfuse上的dataset里面的run就可以看到了，其中有得分，
值越大代表越准确，点击详情还可以看到每一条的评估结果，具体可以项目里面的
评估结果.png和评估明细.png两张图片
'''
